from AStock.ASCrawler import ASStockDayCrawler, ASStockDaySaver
from AStock.ASEvent import Engine, Task
from AStock.ASQuery import ASQuery_stock_info, ASQuery_trade_cal
from AStock.ASUtil import current_date
import argparse
from datetime import datetime


def _crawl_result_handler(data, **kwargs):
    if 'error' in data:
        print(data)
        return

    # print('ASStockDayCrawler result: {}'.format(data))

    event_type = 'stock_day'

    try:
        engine = kwargs['engine']
    except KeyError as e:
        print(e)
        return

    # 保存数据
    event_data = {'event_type': event_type, 'data': data}
    save_task = Task(
        event_data,
        ASStockDaySaver(),
        bind_event='save_stock_day',
        name='task_save',
        callback=print
    )
    engine.run_task(save_task)


def _main_(args):
    cdate = current_date()
    if not ASQuery_trade_cal(cdate):
        print('not trade day, exit...')
        exit(1)
    codes = args.stock_codes
    print('stock codes: {}'.format(codes))
    if codes and isinstance(codes, str):
        codes = [codes]
    elif not codes:
        df = ASQuery_stock_info(fields=['code'])
        codes = df['code'].to_list()

    print('total codes num: {}'.format(len(codes)))

    engine = Engine()
    engine.start()
    engine.create_event('crawl_stock_day')
    engine.start_event('crawl_stock_day')
    engine.create_event('save_stock_day')
    engine.start_event('save_stock_day')

    event_type = 'stock_day'
    for code in codes:
        event_data = {
            'event_type': event_type,
            'code': code
        }
        crawl_task = Task(
            event_data,
            ASStockDayCrawler(),
            bind_event='crawl_stock_day',
            name='task_crawl',
            callback=_crawl_result_handler,
            engine=engine
        )
        engine.run_task(crawl_task)

    # 等待任务完成
    engine.safe_quit()
    print('done')


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='save stock day')
    argparser.add_argument('-c', '--stock-codes', nargs='+', help='stock codes for day price')
    args = argparser.parse_args()
    _main_(args)
